Skip to content

Enable Reactant AD through parallel spin kernels - #922

Open
KookiesNKareem wants to merge 4 commits into
remove-unnecessary-type-paramsfrom
branch/reactant-parallel-ad
Open

Enable Reactant AD through parallel spin kernels#922
KookiesNKareem wants to merge 4 commits into
remove-unnecessary-type-paramsfrom
branch/reactant-parallel-ad

Conversation

@KookiesNKareem

@KookiesNKareem KookiesNKareem commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • bypass ThreadsX and view construction when the parallel excitation/precession wrappers run with one worker, which is the GPU and Reactant path
  • construct the spinor imaginary constant in the input floating type and skip empty precession acquisition work
  • add a focused Reactant + Enzyme reverse-mode regression test against a fifth-order central finite-difference baseline

Validation

  • Reactant loss: 0.6920155690793802 (identical to native)
  • Reactant + Enzyme gradient: [0.019503560659149446, 0.0196742608887282, 0.0]
  • finite-difference gradient: [0.01950356065917826, 0.019674260888777644, -1.1922786455233403e-15]
  • CPU core tests: non-motion 384/384; motion 50/50
  • Metal core tests on Apple M4: non-motion 384/384; motion 50/50
  • final focused Spinor/BlochSimple/AD Metal checks: 69/69

Metal benchmark

BlochSimple, official benchmark fixtures, 30-second BenchmarkTools trials, explicit Metal.synchronize(), two warm runs per source state:

Fixture #814 This PR Average delta
Slice Selection 3D 0.858–0.923 s 0.862–0.891 s -1.5%
MRI Lab 0.715–0.762 s 0.642–0.689 s -10.0%

Stacked on #814.

@KookiesNKareem KookiesNKareem added the run-gpu-ci To run .buildkite CI label Jul 31, 2026 — with ChatGPT Codex Connector
@KookiesNKareem KookiesNKareem mentioned this pull request Jul 31, 2026
Open
@KookiesNKareem
KookiesNKareem marked this pull request as draft August 1, 2026 19:01
@codecov

codecov Bot commented Aug 2, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.65%. Comparing base (2eb6567) to head (9dfc437).

Additional details and impacted files

Impacted file tree graph

@@                       Coverage Diff                       @@
##           remove-unnecessary-type-params     #922   +/-   ##
===============================================================
  Coverage                           91.64%   91.65%           
===============================================================
  Files                                 146      146           
  Lines                                7525     7534    +9     
===============================================================
+ Hits                                 6896     6905    +9     
  Misses                                629      629
Flag Coverage Δ
base 90.96% <ø> (ø)
core 90.19% <100.00%> (+0.06%) ⬆️
files 95.15% <ø> (ø)
komamri 90.51% <ø> (ø)
plots 91.76% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
KomaMRICore/src/datatypes/Spinor.jl 94.44% <100.00%> (+0.69%) ⬆️
...c/simulation/SimMethods/BlochSimple/BlochSimple.jl 89.83% <100.00%> (+0.17%) ⬆️
KomaMRICore/src/simulation/SimulatorCore.jl 92.93% <100.00%> (+0.23%) ⬆️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@KookiesNKareem
KookiesNKareem force-pushed the branch/reactant-parallel-ad branch from c849963 to f2050aa Compare August 2, 2026 15:59
@KookiesNKareem
KookiesNKareem marked this pull request as ready for review August 2, 2026 18:53
cncastillo
cncastillo requested changes Aug 10, 2026
View reviewed changes
Comment thread .buildkite/runtests.yml Outdated
Comment thread KomaMRICore/src/simulation/SimulatorCore.jl Outdated
Comment thread KomaMRICore/src/simulation/SimulatorCore.jl Outdated
Comment thread KomaMRICore/test/test_files/ad_utils.jl Outdated
Comment thread KomaMRICore/test/reactant_cuda.jl Outdated
Comment thread KomaMRICore/test/runtests.jl
github-actions[bot]
github-actions Bot reviewed Aug 10, 2026
View reviewed changes

@github-actions github-actions Bot left a comment •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

KomaMRI Benchmarks

Details
Benchmark suite Current: 3c7280b Previous: 206593e Ratio
MRI Lab/Bloch/GPU/AMDGPU 24431121 ns 21312344 ns 1.15
MRI Lab/Bloch/GPU/CUDA 16415425 ns 17449290 ns 0.94
MRI Lab/Bloch/GPU/Metal 69914209 ns 50797250 ns 1.38
MRI Lab/BlochMagnus2/GPU/AMDGPU 24473021 ns 21285980.5 ns 1.15
MRI Lab/BlochMagnus2/GPU/CUDA 16451149 ns 17455335 ns 0.94
MRI Lab/BlochMagnus2/GPU/Metal 69603250 ns 50882333 ns 1.37
MRI Lab/BlochMagnus4/GPU/AMDGPU 24497981 ns 21389899.5 ns 1.15
MRI Lab/BlochMagnus4/GPU/CUDA 16479139 ns 17553544 ns 0.94
MRI Lab/BlochMagnus4/GPU/Metal 70446854.5 ns 51115958 ns 1.38
Slice Selection 3D/Bloch/GPU/AMDGPU 27497092 ns 27274991.5 ns 1.01
Slice Selection 3D/Bloch/GPU/CUDA 18855924.5 ns 19828203 ns 0.95
Slice Selection 3D/Bloch/GPU/Metal 108584500 ns 79670875 ns 1.36
Slice Selection 3D/BlochMagnus2/GPU/AMDGPU 27505010.5 ns 27256905 ns 1.01
Slice Selection 3D/BlochMagnus2/GPU/CUDA 18917732.5 ns 19761035 ns 0.96
Slice Selection 3D/BlochMagnus2/GPU/Metal 108394666.5 ns 79392250 ns 1.37
Slice Selection 3D/BlochMagnus4/GPU/AMDGPU 27460921.5 ns 27246946.5 ns 1.01
Slice Selection 3D/BlochMagnus4/GPU/CUDA 18845949 ns 19783973 ns 0.95
Slice Selection 3D/BlochMagnus4/GPU/Metal 108849583 ns 79647166 ns 1.37

This comment was automatically generated by workflow using github-action-benchmark.

@KookiesNKareem
KookiesNKareem force-pushed the branch/reactant-parallel-ad branch from 252ac5b to 848e82e Compare August 19, 2026 23:45
@KookiesNKareem
KookiesNKareem force-pushed the branch/reactant-parallel-ad branch from 848e82e to 3c7280b Compare August 21, 2026 04:57
@KookiesNKareem KookiesNKareem mentioned this pull request Aug 22, 2026
Open
@KookiesNKareem
KookiesNKareem force-pushed the branch/reactant-parallel-ad branch from 3c7280b to 9dfc437 Compare August 24, 2026 21:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@github-actions github-actions[bot] github-actions[bot] left review comments
@cncastillo cncastillo cncastillo approved these changes

Assignees

No one assigned

Labels

run-gpu-ci To run .buildkite CI

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@KookiesNKareem @cncastillo